import typing
import os
import struct
import shutil


MOC3_MAGIC = b'MOC3'
MOC3_OFFSETS_V1 = {
    'Counts': 0,
    'Canvas': 1,
    'Parts.space': 2,
    'Parts.id': 3,
    'Deformers.id': 11,
    'ArtMeshes.space1': 29,
    'ArtMeshes.id': 33,
    'ArtMeshes.keyformIndex': 35,
    'ArtMeshes.keyformCount': 36,
    'ArtMeshes.texture': 41,
    'ArtMeshes.drawableVertexCount': 43,
    'ArtMeshes.uvIndexStart': 44,
    'ArtMeshes.drawableIndexStart': 45,
    'ArtMeshes.drawableIndexCount': 46,
    'ArtMeshKeyforms.drawOrder': 69,
    'ArtMeshKeyforms.keyformPositionIndex': 70,
    'KeyformPosition.xy': 71,
    'UVs.xy': 78,
    'DrawableIndexes.index': 79,
}
MOC3_COUNTS_V1 = {
    'Parts': 0,
    'Deformers': 1,
    'ArtMeshes': 4,
    'ArtMeshKeyforms': 9,
    'KeyformPosition': 10,
    'UVs': 15,
    'DrawableIndexes': 16,
}


class Moc3Reader:
    def __init__(self, moc3_dir: str, log_function=print):
        self.moc3_dir = moc3_dir
        self.log_function = log_function
        # read file
        if not os.path.exists(self.moc3_dir):
            raise RuntimeError('file not exists: ' + self.moc3_dir)
        self.fp = open(self.moc3_dir, 'rb+')
        # check header
        header_chunk = self.fp.read(64)
        if header_chunk[:4] != MOC3_MAGIC:
            raise RuntimeError('incorrect magic number')
        self.moc3_version = header_chunk[4]
        if self.moc3_version > 3:
            log_function(f'warning: moc3 version "{self.moc3_version}" is not 100% supported. but may works well.')
        self.moc3_endian = header_chunk[5]
        # mapping
        self._name2offset_idx = MOC3_OFFSETS_V1
        self._name2count_idx = MOC3_COUNTS_V1
        # offset and counts
        self.moc3_offset_list = self._read_chunk(0x40, 0x280)
        self.moc3_counts = self._read_chunk(self.moc3_offset_list[0], 0x80)

        self._chunk_cache = {}

    def save_as(self, file: str, data_uv: typing.Collection[float], data_pos: typing.Collection[float], data_texture_idx: typing.Collection[int]):
        shutil.copy(self.moc3_dir, file)
        self.fp = open(file, 'rb+')
        self.write_back_uv(data_uv)
        self.write_back_pos(data_pos)
        self.write_back_texture_idx(data_texture_idx)

    def _read_chunk_bytes(self, start: int, count: int, unit_size: int):
        self.fp.seek(start)
        return [self.fp.read(unit_size) for _ in range(count)]

    def read_chunk_named_bytes(self, name: str, unit_size: int):
        _idx = name.index(".")
        _count = self.moc3_counts[self._name2count_idx[name[:_idx]]]
        _offset = self.moc3_offset_list[self._name2offset_idx[name]]
        return self._read_chunk_bytes(_offset, _count, unit_size)

    def _read_chunk(self, start: int, size: int, _format='<i'):
        self.fp.seek(start)
        return [x[0] for x in struct.iter_unpack(_format, self.fp.read(size))]

    def read_chunk_named(self, name: str, unit_size=4, _format='<i'):
        _idx = name.index(".")
        _count = self.moc3_counts[self._name2count_idx[name[:_idx]]]
        _offset = self.moc3_offset_list[self._name2offset_idx[name]]
        return self._read_chunk(_offset, _count * unit_size, _format=_format)

    def read_canvas_structure(self) -> typing.Dict[str, float]:
        _offset = self.moc3_offset_list[MOC3_OFFSETS_V1['Canvas']]
        self.fp.seek(_offset)
        _data = [x[0] for x in struct.iter_unpack('<f', self.fp.read(20))]
        result = {'pixelsPerUnit': _data[0], 'originX': _data[1], 'originY': _data[2], 'width': _data[3],
                  'height': _data[4]}
        return result

    def write_back_uv(self, data: typing.Collection[float]):
        self.write_chunk_named(data, 'UVs.xy', _format='<f')

    def write_back_pos(self, data: typing.Collection[float]):
        self.write_chunk_named(data, 'KeyformPosition.xy', _format='<f')

    def write_back_texture_idx(self, data: typing.Collection[int]):
        self.write_chunk_named(data, 'ArtMeshes.texture')

    def _write_chunk(self, data: typing.Iterable, start: int, _format='<i'):
        self.fp.seek(start)
        for d in data:
            self.fp.write(struct.pack(_format, d))

    def write_chunk_named(self, data: typing.Collection, name: str, _format='<i'):
        _idx = name.index(".")
        _count = self.moc3_counts[self._name2count_idx[name[:_idx]]]
        if len(data) > _count:
            print("*WARNING: length of data larger than space cnt")
        _offset = self.moc3_offset_list[self._name2offset_idx[name]]
        self._write_chunk(data, _offset, _format=_format)

    def re_open(self):
        self.fp.flush()
        # re-open file to ensure data is written to disk immediately
        self.fp.close()
        self.fp = open(self.moc3_dir, 'rb+')
